"""Quota and rate limiting management."""

from typing import Dict, Optional, Tuple
from datetime import datetime, timedelta
from dataclasses import dataclass
import asyncio
import motor.motor_asyncio
from bson import ObjectId
import redis.asyncio as redis

from ..models.api_key_model import APIKey, QuotaLimits


@dataclass
class QuotaStatus:
    """Current quota status."""
    daily_tokens_used: int = 0
    daily_tokens_limit: Optional[int] = None
    monthly_tokens_used: int = 0
    monthly_tokens_limit: Optional[int] = None
    cost_used: float = 0.0
    cost_limit: Optional[float] = None
    requests_today: int = 0
    requests_per_day_limit: Optional[int] = None
    is_exceeded: bool = False
    reset_time: Optional[datetime] = None


class QuotaManager:
    """Manages API key quotas and rate limits."""
    
    def __init__(self, mongodb_client: motor.motor_asyncio.AsyncIOMotorClient = None,
                 redis_client: redis.Redis = None):
        """Initialize quota manager."""
        self.mongodb = mongodb_client
        self.redis = redis_client
        
    async def check_quota(self, key_id: str) -> Tuple[bool, Optional[str]]:
        """Check if API key has exceeded quotas.
        
        Args:
            key_id: API key ID
            
        Returns:
            Tuple of (is_allowed, error_message)
        """
        # Get API key from database
        if not self.mongodb:
            return True, None
            
        db = self.mongodb.ai_writing
        doc = await db.api_keys.find_one({"_id": ObjectId(key_id)})
        
        if not doc:
            return False, "API key not found"
            
        api_key = APIKey.from_dict(doc)
        
        # Check if key is active
        if api_key.status.value != "active":
            return False, f"API key is {api_key.status.value}"
            
        # Check expiration
        if api_key.expires_at and api_key.expires_at < datetime.utcnow():
            return False, "API key has expired"
            
        # Get current usage
        status = await self.get_quota_status(key_id)
        
        # Check quotas
        if api_key.quota_limits.daily_tokens and status.daily_tokens_used >= api_key.quota_limits.daily_tokens:
            return False, "Daily token limit exceeded"
            
        if api_key.quota_limits.monthly_tokens and status.monthly_tokens_used >= api_key.quota_limits.monthly_tokens:
            return False, "Monthly token limit exceeded"
            
        if api_key.quota_limits.max_cost and status.cost_used >= api_key.quota_limits.max_cost:
            return False, "Cost limit exceeded"
            
        if api_key.quota_limits.requests_per_day and status.requests_today >= api_key.quota_limits.requests_per_day:
            return False, "Daily request limit exceeded"
            
        # Check rate limiting (requests per minute)
        if api_key.quota_limits.requests_per_minute:
            is_allowed = await self._check_rate_limit(key_id, api_key.quota_limits.requests_per_minute)
            if not is_allowed:
                return False, "Rate limit exceeded (too many requests per minute)"
                
        return True, None
        
    async def _check_rate_limit(self, key_id: str, limit: int) -> bool:
        """Check rate limit using sliding window.
        
        Args:
            key_id: API key ID
            limit: Requests per minute limit
            
        Returns:
            True if request is allowed
        """
        if not self.redis:
            return True
            
        key = f"rate_limit:{key_id}"
        current_minute = int(datetime.utcnow().timestamp() / 60)
        
        pipe = self.redis.pipeline()
        pipe.zadd(key, {str(datetime.utcnow().timestamp()): datetime.utcnow().timestamp()})
        pipe.zremrangebyscore(key, 0, datetime.utcnow().timestamp() - 60)
        pipe.zcard(key)
        pipe.expire(key, 60)
        
        results = await pipe.execute()
        count = results[2]
        
        return count <= limit
        
    async def get_quota_status(self, key_id: str) -> QuotaStatus:
        """Get current quota status for an API key.
        
        Args:
            key_id: API key ID
            
        Returns:
            Quota status
        """
        status = QuotaStatus()
        
        # Try Redis first, fall back to MongoDB if unavailable
        if self.redis:
            try:
                today = datetime.utcnow().strftime("%Y-%m-%d")
                month = datetime.utcnow().strftime("%Y-%m")
                
                daily_key = f"usage:{key_id}:{today}"
                monthly_key = f"usage:{key_id}:{month}"
                
                daily_data = await self.redis.hgetall(daily_key)
                monthly_data = await self.redis.hgetall(monthly_key)
                
                if daily_data:
                    status.daily_tokens_used = int(daily_data.get(b"tokens", 0))
                    status.requests_today = int(daily_data.get(b"requests", 0))
                    
                if monthly_data:
                    status.monthly_tokens_used = int(monthly_data.get(b"tokens", 0))
                    status.cost_used = float(monthly_data.get(b"cost", 0))
                    
            except (redis.ConnectionError, redis.TimeoutError, Exception) as e:
                print(f"Redis unavailable, falling back to MongoDB for quota status: {e}")
                await self._get_quota_status_from_mongodb(key_id, status)
        else:
            # No Redis, use MongoDB directly
            await self._get_quota_status_from_mongodb(key_id, status)
        
        # Get limits from database
        if self.mongodb:
            db = self.mongodb.ai_writing
            doc = await db.api_keys.find_one(
                {"_id": ObjectId(key_id)},
                {"quota_limits": 1}
            )
            
            if doc and "quota_limits" in doc:
                limits = doc["quota_limits"]
                status.daily_tokens_limit = limits.get("daily_tokens")
                status.monthly_tokens_limit = limits.get("monthly_tokens")
                status.cost_limit = limits.get("max_cost")
                status.requests_per_day_limit = limits.get("requests_per_day")
                
        # Calculate reset time
        tomorrow = datetime.utcnow().replace(hour=0, minute=0, second=0, microsecond=0) + timedelta(days=1)
        status.reset_time = tomorrow
        
        # Check if exceeded
        if status.daily_tokens_limit and status.daily_tokens_used >= status.daily_tokens_limit:
            status.is_exceeded = True
        if status.monthly_tokens_limit and status.monthly_tokens_used >= status.monthly_tokens_limit:
            status.is_exceeded = True
        if status.cost_limit and status.cost_used >= status.cost_limit:
            status.is_exceeded = True
            
        return status
        
    async def _get_quota_status_from_mongodb(self, key_id: str, status: QuotaStatus):
        """Fallback method to get quota status from MongoDB.
        
        Args:
            key_id: API key ID
            status: QuotaStatus object to populate
        """
        if not self.mongodb:
            return
            
        try:
            db = self.mongodb.ai_writing
            today = datetime.utcnow().strftime("%Y-%m-%d")
            month = datetime.utcnow().strftime("%Y-%m")
            
            # Get daily stats
            daily_stats = await db.usage_stats.find_one({
                "key_id": key_id,
                "date": today,
                "type": "daily"
            })
            
            if daily_stats:
                status.daily_tokens_used = daily_stats.get("total_tokens", 0)
                status.requests_today = daily_stats.get("request_count", 0)
            
            # Get monthly stats
            monthly_stats = await db.usage_stats.find_one({
                "key_id": key_id,
                "date": month,
                "type": "monthly"
            })
            
            if monthly_stats:
                status.monthly_tokens_used = monthly_stats.get("total_tokens", 0)
                status.cost_used = monthly_stats.get("total_cost", 0.0)
                
        except Exception as e:
            print(f"Failed to get quota status from MongoDB: {e}")
            
    async def _check_rate_limit(self, key_id: str, limit: int) -> bool:
        """Check rate limit using sliding window with Redis fallback.
        
        Args:
            key_id: API key ID
            limit: Requests per minute limit
            
        Returns:
            True if request is allowed
        """
        if not self.redis:
            # Without Redis, we can't do precise rate limiting
            # Fall back to a simple daily request count check
            return await self._check_daily_rate_limit_fallback(key_id, limit)
            
        try:
            key = f"rate_limit:{key_id}"
            current_minute = int(datetime.utcnow().timestamp() / 60)
            
            pipe = self.redis.pipeline()
            pipe.zadd(key, {str(datetime.utcnow().timestamp()): datetime.utcnow().timestamp()})
            pipe.zremrangebyscore(key, 0, datetime.utcnow().timestamp() - 60)
            pipe.zcard(key)
            pipe.expire(key, 60)
            
            results = await pipe.execute()
            count = results[2]
            
            return count <= limit
            
        except (redis.ConnectionError, redis.TimeoutError, Exception) as e:
            print(f"Redis rate limiting failed, using fallback: {e}")
            return await self._check_daily_rate_limit_fallback(key_id, limit)
            
    async def _check_daily_rate_limit_fallback(self, key_id: str, requests_per_minute: int) -> bool:
        """Fallback rate limiting using daily request counts.
        
        Args:
            key_id: API key ID
            requests_per_minute: Requests per minute limit
            
        Returns:
            True if request is allowed (conservative estimate)
        """
        if not self.mongodb:
            return True  # No way to enforce without storage
            
        try:
            # Conservative: allow 60% of theoretical daily max (1440 minutes per day)
            theoretical_daily_max = requests_per_minute * 1440
            conservative_daily_limit = int(theoretical_daily_max * 0.6)
            
            status = QuotaStatus()
            await self._get_quota_status_from_mongodb(key_id, status)
            
            return status.requests_today < conservative_daily_limit
            
        except Exception as e:
            print(f"Daily rate limit fallback failed: {e}")
            return True  # Allow request if we can't determine usage
        
    async def send_quota_alert(self, key_id: str, percentage: int):
        """Send alert when quota reaches certain percentage.
        
        Args:
            key_id: API key ID
            percentage: Current usage percentage
        """
        # TODO: Integrate with notification service
        print(f"Quota alert for key {key_id}: {percentage}% used")


class QuotaMiddleware:
    """Middleware for quota enforcement."""
    
    def __init__(self, quota_manager: QuotaManager):
        """Initialize middleware."""
        self.quota_manager = quota_manager
        
    async def check(self, key_id: str) -> Tuple[bool, Optional[str]]:
        """Check if request is allowed under quotas.
        
        Args:
            key_id: API key ID
            
        Returns:
            Tuple of (is_allowed, error_message)
        """
        return await self.quota_manager.check_quota(key_id)
        
    async def update(self, key_id: str, tokens: int, cost: float):
        """Update usage after successful request.
        
        Args:
            key_id: API key ID
            tokens: Tokens used
            cost: Cost incurred
        """
        # Check if approaching limits and send alerts
        status = await self.quota_manager.get_quota_status(key_id)
        
        if status.daily_tokens_limit:
            usage_percent = (status.daily_tokens_used / status.daily_tokens_limit) * 100
            if usage_percent >= 80 and usage_percent < 100:
                await self.quota_manager.send_quota_alert(key_id, int(usage_percent))